# still, i wanna change this part to database dependent
# maybe sqlalchemy

class Protein(object):
    def __init__(self, id, start, location=-1):
        self.id = id
        self.start = start
        self.location = location
        
def read_gcn_file(file_name):
    id2start = dict()
    taxon_dict = dict() #dict[taxon][contig][strand] = set()
    id2taxon = dict()
    for line in open(file_name):
        try:
            id, start, end, strand, contig, taxon = line.strip().split('\t')[0:6]
        except:
            continue
        id2start[id] = start
        id2taxon[id] = (taxon, contig, strand)
        if taxon not in taxon_dict:
            taxon_dict[taxon] = dict()
        if 'Un' in contig:
            continue
        if contig not in taxon_dict[taxon]:
            taxon_dict[taxon][contig] = {'+':set(), '-':set()}
        if strand == '+' or strand == '-':
            taxon_dict[taxon][contig][strand].add(id)

    return id2start, taxon_dict, id2taxon


def find_location(id2start, taxon_dict):
    id2location = dict()
    for taxon in taxon_dict:
        for contig in taxon_dict[taxon]:
            for strand in taxon_dict[taxon][contig]:
                start2id = dict()
                for id in taxon_dict[taxon][contig][strand]:
                    start2id[int(id2start[id])] = id
                count = 0
                start_list = start2id.keys()
                start_list.sort()
                for start in start_list:
                    count += 1
                    id2location[start2id[start]] = count
                    
                    
def build_profile(orth_file, taxon_id, taxon2id, id2taxon):
    '''del id2taxon after executing this method'''
    id_set = set()
    profile = dict()

    for contig in taxon2id[taxon_id]:
        for strand in taxon2id[taxon_id][contig]:
            id_set = id_set.union(taxon2id[taxon_id][contig][strand])

    for line in open(orth_file):
        tmp = line.strip().split()
        if tmp[0] in id_set or tmp[1] in id_set:
            if tmp[0] in id_set:
                if tmp[0] not in profile:
                    profile[tmp[0]] = {}
                profile[tmp[0]][id2taxon[tmp[1]][0]] = tmp[1]
                continue
            if tmp[1] in id_set:
                if tmp[1] not in profile:
                    profile[tmp[1]] = {}
                profile[tmp[1]][id2taxon[tmp[0]][0]] = tmp[0]

    return profile, id_set

